package util;

import java.util.HashMap;
import java.util.Map;

public class Trie {
	private Map<Character, Trie> childrenMap;
	private boolean isEnd;

	public Trie() {
		childrenMap = new HashMap<>();
		isEnd = false;
	}

	public void insert(String word) {
		Trie node = this;
		for (int i = 0; i < word.length(); i++) {
			char ch = word.charAt(i);
			if (node.childrenMap.get(ch) == null) {
				node.childrenMap.put(ch, new Trie());
			}
			node = node.childrenMap.get(ch);
		}
		node.isEnd = true;
	}

	public boolean search(String word) {
		Trie node = searchPrefix(word);
		return node != null && node.isEnd;
	}

	public boolean startsWith(String prefix) {
		return searchPrefix(prefix) != null;
	}

	private Trie searchPrefix(String prefix) {
		Trie node = this;
		for (int i = 0; i < prefix.length(); i++) {
			char ch = prefix.charAt(i);
			int index = ch - 'a';
			if (node.childrenMap.get(ch) == null) {
				return null;
			}
			node = node.childrenMap.get(ch);
		}
		return node;
	}
}